# -*- coding: utf-8 -*-

import wave, time
from ft.fwth import fwth, ifwth
from ft.fwht import fwht, ifwht
from ft.fwtd import fwtd, ifwtd
from pylab import *

def frames2signal(frames, sampwidth):
    """
    Составляет сигнал из строки фрэймов, которую вернул Wave_read.readframes().
    
    Ожидаемый порядок байтов: little-endian
    """
    
    for i in xrange(0, len(frames), sampwidth):
        # Составляем целое число из фреймов. Некрасиво, зато быстро.
        mostSignificant = i + sampwidth - 1 # позиция старшего байта
        res = ord(frames[mostSignificant])
        for j in xrange(mostSignificant - 1, i - 1, -1):
            res <<= 8
            res |= ord(frames[j]);
        # res --- это и есть значение сигнала
        yield res

def signal2frames(sig, sampwidth):
    """
    Записывает сигнал побайтово в строку. Порядок байтов: little-endian
    """
    
    mask = 0x000000ff
    for num in sig:
        tmp = int(round(num))
        for i in xrange(sampwidth):
            yield chr(tmp & mask)
            tmp >>= 8

def printinfo(wf):
    """
    Выводит в консоль информацию о звуковом файле.
    
    wf: Wave_reader, возвращённый методом wave.open()
    """
    
    if wf.getnchannels() == 1:
        print "mono,",
    else:
        print "stereo,",
    print w.getsampwidth()*8, "bits,",
    print w.getframerate(), "Hz sampling rate"
    print w.getnframes(), "frames"

def getsignal(fn, nframes):
    """
    Берёт сигнал из файла.
    
    fn: имя файла
    nframes: количество фреймов
    """
    
    w = wave.open(fn, "r")
    frames = w.readframes(nframes)
    sig = array(list(frames2signal(frames, w.getsampwidth())))
    w.close()
    return sig

def writesignal(fn, sig, nchannels, framerate, sampwidth):
    """
    Записывает сигнал в файл с заданными настройками.
    
    fn: имя файла
    sig: сигнал
    nchannels: количество каналов
    framerate: частота дискретизации
    sampwidth: количество байт в семпле
    """
    
    w = wave.open(fn, "w")
    w.setnchannels(nchannels)
    w.setframerate(framerate)
    w.setsampwidth(sampwidth)
    frames = ''.join(signal2frames(sig, sampwidth))
    w.writeframes(frames)
    w.close()

def treshold(sp, t):
    counter = 0
    for i in xrange(sp.size):
        if abs(sp[i]) < t:
            sp[i] = 0
            counter += 1
        else:
            sp[i] = sign(sp[i]) * (abs(sp[i]) - t)
    print counter

sig = getsignal("samples/sample1.wav", 1 << 16)
sig /= 1 << 16 # уменьшаем разрешение сигнала
#sinsig = 20000 * sin(linspace(0, 2 * pi, 1 << 8))
#sinsig = [sinsig for i in xrange(1 << 8)]
#sig = concatenate(sinsig)
#sig = linspace(0, 1 << 16 - 1, 1 << 16)
#sig = array([1 << 29], float)
#sig = repeat(sig, 1 << 16)
#sig /= 1 << 16
#sig += 50 * randn(sig.size)
writesignal("samples/noise.wav", sig, 1, 22050, 2)
ordering = "sequency"
subplot(3, 1, 1)
plot(sig)
sp = fwth(sig)
#sp = fwht(sig, ordering)
#treshold(sp, 500.0)
#sp[1 << 15:] = 0
subplot(3, 1, 2)
plot(sp)
resig = ifwth(sp)
#resig = ifwht(sp, ordering)
subplot(3, 1, 3)
plot(resig)
show()
writesignal("samples/sample1o.wav", resig, 1, 22050, 2)
